import socket


class InjectionContext:
    def __init__(self, arg):
        self.client = None
        self.args = arg
        url = arg.dest
        if '://' in url:
            uri_list = url.split("://")
            self.proto = uri_list[0]
            url = uri_list[1]
        self.ip, port = url.split(':')
        self.port = int(port)

    def __del__(self):
        if self.client:
            self.client.close()

    def handle(self, data):
        if self.client is None:
            self.create_client()
        self.send(data)

    def create_client(self):
        if self.proto == "udp":
            self.client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        elif self.proto == "tcp":
            self.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.client.connect((self.ip, self.port))

    def send(self, data):
        if self.proto == "udp":
            self.client.sendto(data, (self.ip, self.port))
        elif self.proto == "tcp":
            self.client.sendall(data)

